package com.yingsheng.bi.hks.template.controller.advice;

import com.yingsheng.bi.hks.template.annotation.Template;
import com.yingsheng.bi.hks.template.controller.IControllerTemplate;
import com.yingsheng.bi.hks.template.events.HandlerMethodEvent;
import com.yingsheng.bi.hks.template.utils.TemplateContext;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.core.MethodParameter;
import org.springframework.core.annotation.AnnotatedElementUtils;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpInputMessage;
import org.springframework.http.MediaType;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.RestControllerAdvice;
import org.springframework.web.servlet.mvc.method.annotation.RequestBodyAdvice;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyAdvice;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.lang.reflect.Type;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@RestControllerAdvice
public class RequestResponseBodyAdvice implements RequestBodyAdvice, ResponseBodyAdvice, ApplicationContextAware, InitializingBean {

    private final Logger log = LoggerFactory.getLogger(RequestResponseBodyAdvice.class);

    public static final Charset DEFAULT_CHARSET = StandardCharsets.UTF_8;

    private final Map<String, IControllerTemplate> templateMap = new ConcurrentHashMap<>();

    private ApplicationContext applicationContext;

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
    }

    public IControllerTemplate getTemplate(String name) {
        if (StringUtils.hasText(name)) {
            return this.templateMap.get(name);
        }
        return null;
    }

    @Override
    public boolean supports(MethodParameter methodParameter, Type targetType, Class<? extends HttpMessageConverter<?>> converterType) {
        Template templateAnnotation = AnnotatedElementUtils.findMergedAnnotation(methodParameter.getMethod(), Template.class);
        if (templateAnnotation == null || !StringUtils.hasText(templateAnnotation.name())) {
            templateAnnotation = AnnotatedElementUtils.findMergedAnnotation(methodParameter.getDeclaringClass(), Template.class);
        }
        if (templateAnnotation != null) {
            IControllerTemplate templateService = getTemplate(templateAnnotation.name());
            if (templateService != null) {
                IControllerTemplate parent = getTemplate(templateAnnotation.parent());
                templateService.setParent(parent);
                TemplateContext templateContext = TemplateContext.getRequestContext();
                templateContext.setTemplate(templateAnnotation.name());
                templateContext.setSupport(true);
                return true;
            }
        }
        return false;
    }

    @Override
    public HttpInputMessage beforeBodyRead(HttpInputMessage inputMessage, MethodParameter parameter, Type targetType, Class<? extends HttpMessageConverter<?>> converterType) throws IOException {
        HandlerMethodEvent event = new HandlerMethodEvent();
        try {
            Charset charset = getCharset(inputMessage);
            event.setInput(inputMessage);
            event.setCharset(charset);
            event.setMethodParameter(parameter);
            TemplateContext templateContext = TemplateContext.getRequestContext();
            templateContext.setEvent(event);
            String template = templateContext.getTemplate();
            IControllerTemplate templateService = getTemplate(template);
            templateService.init(event);

            if (event.getOutput() != null) {
                if (event.getOutput() instanceof HttpInputMessage) {
                    return (HttpInputMessage) event.getOutput();
                } else {
                    Object output = event.getOutput();
                    return new HttpInputMessage() {
                        @Override
                        public HttpHeaders getHeaders() {
                            return inputMessage.getHeaders();
                        }

                        @Override
                        public InputStream getBody() {
                            if (output instanceof String) {
                                byte[] bodyByteArray = ((String) output).getBytes(charset);
                                return new ByteArrayInputStream(bodyByteArray);
                            } else if (output instanceof byte[]) {
                                return new ByteArrayInputStream((byte[]) output);
                            } else {
                                byte[] bodyByteArray = output.toString().getBytes(charset);
                                return new ByteArrayInputStream(bodyByteArray);
                            }
                        }
                    };
                }
            } else {
                return inputMessage;
            }
        } finally {
            event.reset();
        }
    }

    @Override
    public Object afterBodyRead(Object body, HttpInputMessage inputMessage, MethodParameter parameter, Type targetType, Class<? extends HttpMessageConverter<?>> converterType) {
        TemplateContext templateContext = TemplateContext.getRequestContext();
        HandlerMethodEvent event = (HandlerMethodEvent) templateContext.getEvent();
        try {
            String template = templateContext.getTemplate();
            IControllerTemplate templateService = getTemplate(template);
            event.setInput(body);
            event.setRequest(body);
            event.setMethodParameter(parameter);
            templateService.entry(event);
            Class responseClazz = body.getClass();
            if (event.getOutput() != null && event.getOutput().getClass().isAssignableFrom(responseClazz)) {
                return event.getOutput();
            } else {
                return body;
            }
        } finally {
            event.reset();
        }
    }

    @Override
    public Object handleEmptyBody(Object body, HttpInputMessage inputMessage, MethodParameter parameter, Type targetType, Class<? extends HttpMessageConverter<?>> converterType) {
        return body;
    }

    @Override
    public boolean supports(MethodParameter returnType, Class converterType) {
        return TemplateContext.getRequestContext().isSupport();
    }

    @Override
    public Object beforeBodyWrite(Object body, MethodParameter returnType, MediaType selectedContentType, Class selectedConverterType, ServerHttpRequest request, ServerHttpResponse response) {
        try {
            TemplateContext templateContext = TemplateContext.getRequestContext();
            String template = templateContext.getTemplate();
            IControllerTemplate templateService = getTemplate(template);
            HandlerMethodEvent event = (HandlerMethodEvent) templateContext.getEvent();
            event.setMethodParameter(returnType);
            event.setInput(body);
            event.setResponse(body);
            templateService.exit(event);
            if (event.getOutput() != null) {
                return event.getOutput();
            } else if (event.getResponse() != null) {
                return event.getResponse();
            } else {
                return body;
            }
        } finally {
            TemplateContext.reset();
        }
    }

    @Override
    public void afterPropertiesSet() {
        try {
            Map<String, IControllerTemplate> templateMap = this.applicationContext.getBeansOfType(IControllerTemplate.class);
            this.templateMap.putAll(templateMap);
            if (log.isDebugEnabled()) {
                log.debug("Controller Templates: " + this.templateMap.keySet());
            }
        } catch (BeansException e) {
            log.warn(e.getMessage());
        }
    }

    private Charset getCharset(HttpInputMessage inputMessage) {
        HttpHeaders headers = inputMessage.getHeaders();
        Charset charset = headers.getContentType() != null ? headers.getContentType().getCharset() : null;
        return charset != null ? charset : DEFAULT_CHARSET;
    }
}
